use pgls_analyse::{Rule, RuleDiagnostic, RuleSource, context::RuleContext, declare_lint_rule};
use pgls_console::markup;
use pgls_diagnostics::Severity;

declare_lint_rule! {
    /// Prefer using IDENTITY columns over serial columns.
    ///
    /// SERIAL types (serial, serial2, serial4, serial8, smallserial, bigserial) use sequences behind
    /// the scenes but with some limitations. IDENTITY columns provide better control over sequence
    /// behavior and are part of the SQL standard.
    ///
    /// IDENTITY columns offer clearer ownership semantics - the sequence is directly tied to the column
    /// and will be automatically dropped when the column or table is dropped. They also provide better
    /// control through GENERATED ALWAYS (prevents manual inserts) or GENERATED BY DEFAULT options.
    ///
    /// ## Examples
    ///
    /// ### Invalid
    ///
    /// ```sql,expect_diagnostic
    /// create table users (
    ///     id serial
    /// );
    /// ```
    ///
    /// ```sql,expect_diagnostic
    /// create table users (
    ///     id bigserial
    /// );
    /// ```
    ///
    /// ### Valid
    ///
    /// ```sql
    /// create table users (
    ///     id bigint generated by default as identity primary key
    /// );
    /// ```
    ///
    /// ```sql
    /// create table users (
    ///     id bigint generated always as identity primary key
    /// );
    /// ```
    ///
    pub PreferIdentity {
        version: "next",
        name: "preferIdentity",
        severity: Severity::Warning,
        recommended: false,
        sources: &[RuleSource::Squawk("prefer-identity")],
    }
}

impl Rule for PreferIdentity {
    type Options = ();

    fn run(ctx: &RuleContext<Self>) -> Vec<RuleDiagnostic> {
        let mut diagnostics = Vec::new();

        match ctx.stmt() {
            pgls_query::NodeEnum::CreateStmt(stmt) => {
                for table_elt in &stmt.table_elts {
                    if let Some(pgls_query::NodeEnum::ColumnDef(col_def)) = &table_elt.node {
                        check_column_def(&mut diagnostics, col_def);
                    }
                }
            }
            pgls_query::NodeEnum::AlterTableStmt(stmt) => {
                for cmd in &stmt.cmds {
                    if let Some(pgls_query::NodeEnum::AlterTableCmd(cmd)) = &cmd.node {
                        if matches!(
                            cmd.subtype(),
                            pgls_query::protobuf::AlterTableType::AtAddColumn
                                | pgls_query::protobuf::AlterTableType::AtAlterColumnType
                        ) {
                            if let Some(pgls_query::NodeEnum::ColumnDef(col_def)) =
                                &cmd.def.as_ref().and_then(|d| d.node.as_ref())
                            {
                                check_column_def(&mut diagnostics, col_def);
                            }
                        }
                    }
                }
            }
            _ => {}
        }

        diagnostics
    }
}

fn check_column_def(
    diagnostics: &mut Vec<RuleDiagnostic>,
    col_def: &pgls_query::protobuf::ColumnDef,
) {
    let Some(type_name) = &col_def.type_name else {
        return;
    };

    for name_node in &type_name.names {
        let Some(pgls_query::NodeEnum::String(name)) = &name_node.node else {
            continue;
        };

        if !matches!(
            name.sval.as_str(),
            "serial" | "serial2" | "serial4" | "serial8" | "smallserial" | "bigserial"
        ) {
            continue;
        }

        diagnostics.push(
            RuleDiagnostic::new(
                rule_category!(),
                None,
                markup! {
                    "Prefer IDENTITY columns over SERIAL types."
                },
            )
            .detail(None, format!("Column uses '{}' type which has limitations compared to IDENTITY columns.", name.sval))
            .note("Use 'bigint GENERATED BY DEFAULT AS IDENTITY' or 'bigint GENERATED ALWAYS AS IDENTITY' instead."),
        );
    }
}
